#ifndef PICCOLO_VK_PIPELINE_H
#define PICCOLO_VK_PIPELINE_H

#include "PiccoloVKCommon.h"

namespace Piccolo
{
    class PiccoloVKDevice;
    class PiccoloVKRenderPass;

    struct ShaderLayout
    {
        std::vector<VkDescriptorSetLayout> descriptorSetLayouts;
        std::vector<VkPushConstantRange> pushConstantRanges;
    };

    struct PipelineVertexInputState
    {
        std::vector<VkVertexInputBindingDescription> vertexBindingDescriptions;
        std::vector<VkVertexInputAttributeDescription> vertexAttributes;
    };

    struct PipelineInputAssemblyState
    {
        VkPrimitiveTopology topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
        VkBool32 primitiveRestartEnable = VK_FALSE;
    };

    struct PipelineRasterizationState
    {
        VkBool32 depthClampEnable = VK_FALSE;
        VkBool32 rasterizerDiscardEnable = VK_FALSE;
        VkPolygonMode polygonMode = VK_POLYGON_MODE_FILL;
        VkCullModeFlags cullMode = VK_CULL_MODE_NONE;
        VkFrontFace frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
        VkBool32 depthBiasEnable = VK_FALSE;
        float depthBiasConstantFactor = 0;
        float depthBiasClamp = 0;
        float depthBiasSlopeFactor = 0;
        float lineWidth = 1.f;
    };

    struct PipelineMultisampleState
    {
        VkSampleCountFlagBits rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
        VkBool32 sampleShadingEnable = VK_FALSE;
        float minSampleShading = 0.2f;
    };

    struct PipelineDepthStencilState
    {
        VkBool32 depthTestEnable = VK_FALSE;
        VkBool32 depthWriteEnable = VK_FALSE;
        VkCompareOp depthCompareOp = VK_COMPARE_OP_NEVER;
        VkBool32 depthBoundsTestEnable = VK_FALSE;
        VkBool32 stencilTestEnable = VK_FALSE;
    };

    struct PipelineDynamicState
    {
        std::vector<VkDynamicState> dynamicStates;
    };

    struct PipelineConfig
    {
        PipelineVertexInputState vertexInputState;
        PipelineInputAssemblyState inputAssemblyState;
        PipelineRasterizationState rasterizationState;
        PipelineMultisampleState multisampleState;
        PipelineDepthStencilState depthStencilState;
        VkPipelineColorBlendAttachmentState colorBlendAttachmentState
        {
            .blendEnable = VK_FALSE,
            .srcColorBlendFactor = VK_BLEND_FACTOR_ONE,
            .dstColorBlendFactor = VK_BLEND_FACTOR_ZERO,
            .colorBlendOp = VK_BLEND_OP_ADD,
            .srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE,
            .dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
            .alphaBlendOp = VK_BLEND_OP_ADD,
            .colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT
        };
        PipelineDynamicState dynamicState;
    };

    class PiccoloVKPipelineLayout
    {
    public:
        PiccoloVKPipelineLayout(PiccoloVKDevice* device, const std::string& vertexShaderFile, const std::string& fragShaderFile, const ShaderLayout& shaderLayout = {});
        ~PiccoloVKPipelineLayout();

        VkPipelineLayout GetPipelineLayout() const { return mPipelineLayout; }
        VkShaderModule GetVertexShaderModule() const { return mVertexShaderModule; }
        VkShaderModule GetFragShaderModule() const { return mFragShaderModule; }

    private:
        VkResult CreateShaderModule(const std::string& shaderFile, VkShaderModule* shaderModule);

        VkPipelineLayout mPipelineLayout = VK_NULL_HANDLE;

        PiccoloVKDevice* mDevice;
        VkShaderModule mVertexShaderModule = VK_NULL_HANDLE;
        VkShaderModule mFragShaderModule = VK_NULL_HANDLE;
    };

    class PiccoloVKPipeline
    {
    public:
        PiccoloVKPipeline(PiccoloVKDevice* device, PiccoloVKRenderPass* renderPass, PiccoloVKPipelineLayout* pipelineLayout);
        ~PiccoloVKPipeline();

        void Create();
        void Bind(VkCommandBuffer commandBuffer);

        PiccoloVKPipeline* SetVertexInputState(const std::vector<VkVertexInputBindingDescription>& vertexBindings, const std::vector<VkVertexInputAttributeDescription>& vertexAttributes);
        PiccoloVKPipeline* SetInputAssemblyState(VkPrimitiveTopology topology, VkBool32 primitiveRestartEnable = VK_FALSE);
        PiccoloVKPipeline* SetRasterizationState(const PipelineRasterizationState& rasterizationState);
        PiccoloVKPipeline* SetMultisampleState(VkSampleCountFlagBits samples, VkBool32 sampleShadingEnable, float minSampleShading = 0.f);
        PiccoloVKPipeline* SetDepthStencilState(const PipelineDepthStencilState& depthStencilState);
        PiccoloVKPipeline* SetColorBlendAttachmentState(VkBool32 blendEnable,
            VkBlendFactor srcColorBlendFactor = VK_BLEND_FACTOR_ONE, VkBlendFactor dstColorBlendFactor = VK_BLEND_FACTOR_ZERO, VkBlendOp colorBlendOp = VK_BLEND_OP_ADD,
            VkBlendFactor srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE, VkBlendFactor dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO, VkBlendOp alphaBlendOp = VK_BLEND_OP_ADD);
        PiccoloVKPipeline* SetDynamicState(const std::vector<VkDynamicState>& dynamicStates);
        PiccoloVKPipeline* EnableAlphaBlend();
        PiccoloVKPipeline* EnableDepthTest();

        VkPipeline GetPipeline() const { return mPipeline; }

    private:
        VkPipeline mPipeline = VK_NULL_HANDLE;

        PiccoloVKDevice* mDevice;
        PipelineConfig mPipelineConfig;
        PiccoloVKRenderPass *mRenderPass;
        PiccoloVKPipelineLayout *mPipelineLayout;
    };
}

#endif